import os

import pandas as pd
from sqlalchemy import create_engine
import yaml
from General.ConfigLoader import load_config

# 需要安装 pip install mysql-connector-python

class DatabaseReader:
    """
    数据库读取类
    """
    def __init__(self, config_file='config.yaml'):
        """
        初始化 DatabaseReader 类，创建数据库连接引擎。
        """
        config = load_config()
        mysql_config = config['mysql']
        self.engine = create_engine(
            f"mysql+mysqlconnector://{mysql_config['user']}:{mysql_config['password']}@{mysql_config['host']}/{mysql_config['database']}"
        )

    def read_table_by_sql(self, query):
        """
        从数据库中读取数据表。
        :param query: SQL 查询语句，用于从数据库中读取数据。
        :return: 包含查询结果的 Pandas DataFrame。
        """
        return pd.read_sql(query, con=self.engine)

    def read_robot_node(self):
        """
        读取robot_node表
        :return:
        """
        return self.read_table_by_sql("SELECT * FROM robot_node WHERE is_deleted = 0")

    def read_robot_node_lift(self):
        """
        读取robot_node_lift表
        :return:
        """
        return self.read_table_by_sql("SELECT * FROM robot_node_lift")

    def read_robot_edge(self):
        """
        读取robot_edge表
        :return:
        """
        return self.read_table_by_sql("SELECT * FROM robot_edge WHERE is_deleted = 0")

    def read_robot_order(self):
        """
        读取robot_order表
        :return:
        """
        return self.read_table_by_sql("SELECT create_time, start_time, put_time, put_finish_time, get_time, finish_time, start_node, end_node, takelength, taketime "
                                      "FROM robot_order WHERE order_state = 5 AND start_node IS NOT NULL AND end_node IS NOT NULL AND start_time IS NOT NULL AND put_time IS NOT NULL AND get_time IS NOT NULL")

    def read_robot_order_mini(self):
        """
        读取robot_order表，读取1w条小规模数据
        :return:
        """
        return self.read_table_by_sql("SELECT create_time, start_time, put_time, put_finish_time, get_time, finish_time, start_node, end_node, takelength, taketime "
                                      "FROM robot_order WHERE order_state = 5 AND start_node IS NOT NULL AND end_node IS NOT NULL AND start_time IS NOT NULL AND put_time IS NOT NULL AND get_time IS NOT NULL LIMIT 10000")



if __name__ == '__main__':
    # test
    db_reader = DatabaseReader()
    print(db_reader.read_robot_node())